from flask import Flask, render_template, request, redirect, url_for, flash, jsonify
import sqlite3
import datetime
import os
import csv
from io import StringIO

app = Flask(__name__)
app.secret_key = 'your_secret_key'  # 用于flash消息

# 初始化数据库
def init_db():
    conn = sqlite3.connect('finance.db')
    cursor = conn.cursor()
    
    # 创建交易表
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS transactions (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            type TEXT NOT NULL,  -- 'income' 或 'expense'
            amount REAL NOT NULL,
            category TEXT NOT NULL,
            date TEXT NOT NULL,
            description TEXT
        )
    ''')
    
    # 创建类别表
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS categories (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            name TEXT NOT NULL,
            type TEXT NOT NULL,  -- 'income' 或 'expense'
            UNIQUE(name, type)
        )
    ''')
    
    # 创建预算表
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS budgets (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            category TEXT NOT NULL,
            amount REAL NOT NULL,
            month TEXT NOT NULL,
            UNIQUE(category, month)
        )
    ''')
    
    # 插入默认类别
    default_categories = [
        # 收入类别
        ('工资', 'income'),
        ('奖金', 'income'),
        ('投资收益', 'income'),
        ('兼职', 'income'),
        ('其他收入', 'income'),
        # 支出类别
        ('餐饮', 'expense'),
        ('交通', 'expense'),
        ('住房', 'expense'),
        ('购物', 'expense'),
        ('娱乐', 'expense'),
        ('医疗', 'expense'),
        ('教育', 'expense'),
        ('其他支出', 'expense')
    ]
    
    try:
        cursor.executemany('INSERT INTO categories (name, type) VALUES (?, ?)', default_categories)
    except sqlite3.IntegrityError:
        # 类别已存在时忽略
        pass
    
    # 插入一些示例交易数据（如果表为空）
    cursor.execute('SELECT COUNT(*) FROM transactions')
    count = cursor.fetchone()[0]
    
    if count == 0:
        now = datetime.datetime.now()
        sample_transactions = [
            ('income', 8000.0, '工资', (now - datetime.timedelta(days=10)).strftime('%Y-%m-%d'), '月工资'),
            ('expense', 1500.0, '住房', (now - datetime.timedelta(days=8)).strftime('%Y-%m-%d'), '房租'),
            ('expense', 800.0, '餐饮', (now - datetime.timedelta(days=5)).strftime('%Y-%m-%d'), '日常餐饮'),
            ('expense', 300.0, '交通', (now - datetime.timedelta(days=3)).strftime('%Y-%m-%d'), '公交地铁'),
            ('income', 2000.0, '兼职', (now - datetime.timedelta(days=2)).strftime('%Y-%m-%d'), '周末兼职'),
            ('expense', 500.0, '购物', now.strftime('%Y-%m-%d'), '日用品')
        ]
        
        cursor.executemany('INSERT INTO transactions (type, amount, category, date, description) VALUES (?, ?, ?, ?, ?)', sample_transactions)
    
    conn.commit()
    conn.close()

# 获取所有类别
def get_categories(type=None):
    conn = sqlite3.connect('finance.db')
    cursor = conn.cursor()
    
    if type:
        cursor.execute('SELECT name FROM categories WHERE type = ? ORDER BY name', (type,))
    else:
        cursor.execute('SELECT name, type FROM categories ORDER BY type, name')
    
    categories = cursor.fetchall()
    conn.close()
    
    if type:
        # 只返回类别名称的列表
        return [category[0] for category in categories]
    else:
        # 返回包含名称和类型的列表
        return categories

# 添加新交易
def add_transaction(transaction_type, amount, category, date, description):
    conn = sqlite3.connect('finance.db')
    cursor = conn.cursor()
    cursor.execute('INSERT INTO transactions (type, amount, category, date, description) VALUES (?, ?, ?, ?, ?)', 
                  (transaction_type, amount, category, date, description))
    transaction_id = cursor.lastrowid
    conn.commit()
    conn.close()
    return transaction_id

# 获取交易列表
def get_transactions(filters=None):
    conn = sqlite3.connect('finance.db')
    cursor = conn.cursor()
    
    query = 'SELECT id, type, amount, category, date, description FROM transactions WHERE 1=1'
    params = []
    
    if filters:
        if 'type' in filters and filters['type']:
            query += ' AND type = ?'
            params.append(filters['type'])
        
        if 'category' in filters and filters['category']:
            query += ' AND category = ?'
            params.append(filters['category'])
        
        if 'start_date' in filters and filters['start_date']:
            query += ' AND date >= ?'
            params.append(filters['start_date'])
        
        if 'end_date' in filters and filters['end_date']:
            query += ' AND date <= ?'
            params.append(filters['end_date'])
    
    query += ' ORDER BY date DESC'
    
    cursor.execute(query, params)
    transactions = cursor.fetchall()
    conn.close()
    return transactions

# 获取交易详情
def get_transaction_by_id(transaction_id):
    conn = sqlite3.connect('finance.db')
    cursor = conn.cursor()
    cursor.execute('SELECT id, type, amount, category, date, description FROM transactions WHERE id = ?', (transaction_id,))
    transaction = cursor.fetchone()
    conn.close()
    return transaction

# 更新交易
def update_transaction(transaction_id, transaction_type, amount, category, date, description):
    conn = sqlite3.connect('finance.db')
    cursor = conn.cursor()
    cursor.execute('UPDATE transactions SET type = ?, amount = ?, category = ?, date = ?, description = ? WHERE id = ?', 
                  (transaction_type, amount, category, date, description, transaction_id))
    conn.commit()
    conn.close()

# 删除交易
def delete_transaction(transaction_id):
    conn = sqlite3.connect('finance.db')
    cursor = conn.cursor()
    cursor.execute('DELETE FROM transactions WHERE id = ?', (transaction_id,))
    conn.commit()
    conn.close()

# 获取财务统计数据
def get_financial_summary(start_date=None, end_date=None):
    conn = sqlite3.connect('finance.db')
    cursor = conn.cursor()
    
    # 获取总收入
    income_query = 'SELECT SUM(amount) FROM transactions WHERE type = "income"'
    expense_query = 'SELECT SUM(amount) FROM transactions WHERE type = "expense"'
    params = []
    
    if start_date or end_date:
        date_filter = ' AND date BETWEEN ? AND ?'
        params = [start_date or '1900-01-01', end_date or '2100-12-31']
        income_query += date_filter
        expense_query += date_filter
    
    cursor.execute(income_query, params)
    total_income = cursor.fetchone()[0] or 0
    
    cursor.execute(expense_query, params)
    total_expense = cursor.fetchone()[0] or 0
    
    # 获取各类别支出统计
    category_expense_query = 'SELECT category, SUM(amount) FROM transactions WHERE type = "expense"'
    if start_date or end_date:
        category_expense_query += date_filter
    category_expense_query += ' GROUP BY category ORDER BY SUM(amount) DESC'
    
    cursor.execute(category_expense_query, params)
    category_expenses = cursor.fetchall()
    
    conn.close()
    
    return {
        'total_income': total_income,
        'total_expense': total_expense,
        'balance': total_income - total_expense,
        'category_expenses': category_expenses
    }

# 设置预算
def set_budget(category, amount, month):
    conn = sqlite3.connect('finance.db')
    cursor = conn.cursor()
    
    # 检查是否已存在该类别的预算
    cursor.execute('SELECT id FROM budgets WHERE category = ? AND month = ?', (category, month))
    existing_budget = cursor.fetchone()
    
    if existing_budget:
        # 更新现有预算
        cursor.execute('UPDATE budgets SET amount = ? WHERE id = ?', (amount, existing_budget[0]))
    else:
        # 创建新预算
        cursor.execute('INSERT INTO budgets (category, amount, month) VALUES (?, ?, ?)', (category, amount, month))
    
    conn.commit()
    conn.close()

# 获取预算

def get_budgets(month=None):
    conn = sqlite3.connect('finance.db')
    cursor = conn.cursor()
    
    if month:
        cursor.execute('SELECT category, amount FROM budgets WHERE month = ?', (month,))
    else:
        cursor.execute('SELECT category, amount, month FROM budgets ORDER BY month DESC')
    
    budgets = cursor.fetchall()
    conn.close()
    return budgets

# 获取月度预算使用情况
def get_budget_status(month):
    conn = sqlite3.connect('finance.db')
    cursor = conn.cursor()
    
    # 获取该月的所有预算
    cursor.execute('SELECT category, amount FROM budgets WHERE month = ?', (month,))
    budgets = cursor.fetchall()
    
    # 获取该月各类别的实际支出
    status = []
    for category, budget_amount in budgets:
        cursor.execute('SELECT SUM(amount) FROM transactions WHERE type = "expense" AND category = ? AND date LIKE ?', 
                      (category, f'{month}%'))
        actual_amount = cursor.fetchone()[0] or 0
        
        status.append({
            'category': category,
            'budget_amount': budget_amount,
            'actual_amount': actual_amount,
            'remaining': budget_amount - actual_amount,
            'percentage': (actual_amount / budget_amount) * 100 if budget_amount > 0 else 0
        })
    
    conn.close()
    return status

# 主页路由\@app.route('/')
def index():
    # 获取筛选条件
    transaction_type = request.args.get('type', '')
    category = request.args.get('category', '')
    start_date = request.args.get('start_date', '')
    end_date = request.args.get('end_date', '')
    
    # 构建筛选参数
    filters = {}
    if transaction_type:
        filters['type'] = transaction_type
    if category:
        filters['category'] = category
    if start_date:
        filters['start_date'] = start_date
    if end_date:
        filters['end_date'] = end_date
    
    # 获取交易列表
    transactions = get_transactions(filters)
    
    # 获取收入和支出类别
    income_categories = get_categories('income')
    expense_categories = get_categories('expense')
    
    # 获取财务摘要
    summary = get_financial_summary(start_date, end_date)
    
    return render_template('index.html', 
                          transactions=transactions, 
                          income_categories=income_categories, 
                          expense_categories=expense_categories, 
                          summary=summary, 
                          filters={'type': transaction_type, 'category': category, 'start_date': start_date, 'end_date': end_date})

# 添加交易路由\@app.route('/add', methods=['GET', 'POST'])
def add_transaction_route():
    if request.method == 'POST':
        transaction_type = request.form['type']
        amount = float(request.form['amount'])
        category = request.form['category']
        date = request.form['date']
        description = request.form['description'] or None
        
        if amount <= 0:
            flash('金额必须大于0！')
            return redirect(url_for('add_transaction_route'))
        
        add_transaction(transaction_type, amount, category, date, description)
        flash('交易添加成功！')
        return redirect(url_for('index'))
    
    # 获取收入和支出类别
    income_categories = get_categories('income')
    expense_categories = get_categories('expense')
    
    # 设置默认日期为今天
    today = datetime.date.today().strftime('%Y-%m-%d')
    
    return render_template('add_transaction.html', 
                          income_categories=income_categories, 
                          expense_categories=expense_categories, 
                          today=today)

# 编辑交易路由\@app.route('/edit/<int:transaction_id>', methods=['GET', 'POST'])
def edit_transaction_route(transaction_id):
    transaction = get_transaction_by_id(transaction_id)
    
    if not transaction:
        flash('交易不存在！')
        return redirect(url_for('index'))
    
    if request.method == 'POST':
        transaction_type = request.form['type']
        amount = float(request.form['amount'])
        category = request.form['category']
        date = request.form['date']
        description = request.form['description'] or None
        
        if amount <= 0:
            flash('金额必须大于0！')
            return redirect(url_for('edit_transaction_route', transaction_id=transaction_id))
        
        update_transaction(transaction_id, transaction_type, amount, category, date, description)
        flash('交易更新成功！')
        return redirect(url_for('index'))
    
    # 获取收入和支出类别
    income_categories = get_categories('income')
    expense_categories = get_categories('expense')
    
    return render_template('add_transaction.html', 
                          transaction=transaction, 
                          income_categories=income_categories, 
                          expense_categories=expense_categories, 
                          today='')

# 删除交易路由\@app.route('/delete/<int:transaction_id>')
def delete_transaction_route(transaction_id):
    transaction = get_transaction_by_id(transaction_id)
    
    if not transaction:
        flash('交易不存在！')
    else:
        delete_transaction(transaction_id)
        flash('交易删除成功！')
    
    return redirect(url_for('index'))

# 财务报表路由\@app.route('/reports')
def reports():
    # 获取当前月份
    current_month = datetime.datetime.now().strftime('%Y-%m')
    selected_month = request.args.get('month', current_month)
    
    # 构建该月的开始和结束日期
    year, month = selected_month.split('-')
    start_date = f'{year}-{month}-01'
    # 获取该月的最后一天
    if int(month) == 12:
        next_month = f'{int(year)+1}-01'
    else:
        next_month = f'{year}-{int(month)+1:02d}'
    end_date = (datetime.datetime.strptime(next_month, '%Y-%m') - datetime.timedelta(days=1)).strftime('%Y-%m-%d')
    
    # 获取该月的财务统计
    summary = get_financial_summary(start_date, end_date)
    
    # 准备图表数据
    categories = [category for category, _ in summary['category_expenses']]
    amounts = [amount for _, amount in summary['category_expenses']]
    
    # 获取预算状态
    budget_status = get_budget_status(selected_month)
    
    return render_template('reports.html', 
                          summary=summary, 
                          selected_month=selected_month,
                          categories=categories,
                          amounts=amounts,
                          budget_status=budget_status)

# 预算管理路由\@app.route('/budget', methods=['GET', 'POST'])
def budget():
    # 获取当前月份
    current_month = datetime.datetime.now().strftime('%Y-%m')
    selected_month = request.args.get('month', current_month)
    
    if request.method == 'POST':
        category = request.form['category']
        amount = float(request.form['amount'])
        month = request.form['month']
        
        if amount < 0:
            flash('预算金额不能为负数！')
            return redirect(url_for('budget', month=month))
        
        set_budget(category, amount, month)
        flash('预算设置成功！')
        return redirect(url_for('budget', month=month))
    
    # 获取支出类别
    expense_categories = get_categories('expense')
    
    # 获取该月的预算
    budgets = dict(get_budgets(selected_month))
    
    # 获取预算状态
    budget_status = get_budget_status(selected_month)
    
    return render_template('budget.html', 
                          expense_categories=expense_categories,
                          selected_month=selected_month,
                          budgets=budgets,
                          budget_status=budget_status)

# 设置页面路由\@app.route('/settings')
def settings():
    return render_template('settings.html')

# 导出数据API\@app.route('/export')
def export_data():
    # 获取所有交易数据
    transactions = get_transactions()
    
    # 创建CSV文件
    output = StringIO()
    writer = csv.writer(output)
    
    # 写入表头
    writer.writerow(['类型', '金额', '类别', '日期', '描述'])
    
    # 写入数据
    for transaction in transactions:
        transaction_type = '收入' if transaction[1] == 'income' else '支出'
        writer.writerow([transaction_type, transaction[2], transaction[3], transaction[4], transaction[5] or ''])
    
    # 返回CSV响应
    return app.response_class(
        response=output.getvalue(),
        mimetype='text/csv',
        headers={'Content-Disposition': 'attachment; filename=financial_data.csv'}
    )

# 启动时初始化数据库
if __name__ == '__main__':
    init_db()
    print('个人财务管理系统已启动！访问 http://localhost:5000 开始使用。')
    app.run(debug=True)